from typing import Optional, List
from pymilvus import connections, utility


# ======================================================
#  linkMilvus —— 数据库连接层
# ======================================================

class linkMilvus:
    def __init__(
        self,
        alias: str = "default",
        host: str = "localhost",
        port: int = 19530,
        user: Optional[str] = None,
        password: Optional[str] = None,
        auto_connect: bool = True,  # 是否自动连接
    ):
        self.alias = alias
        self.host = host
        self.port = int(port)
        self.user = user
        self.password = password

        if auto_connect:
            self._connect()

    def _connect(self):
        try:
            if connections.has_connection(self.alias):
                connections.disconnect(self.alias)

            params = {"alias": self.alias,
                      "host": self.host, "port": int(self.port)}
            if self.user and self.password:
                params.update(user=self.user, password=self.password)
            connections.connect(**params)
            print(f"已连接 Milvus ({self.host}:{self.port}, alias={self.alias})")
        except Exception as e:
            print(f"[Error] 无法连接 Milvus: {e}")

    def close(self):
        connections.disconnect(self.alias)
        print(f"连接 '{self.alias}' 已断开")

    # ------------------- 上下文管理 -------------------
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    # ------------------- 集合管理 -------------------
    def dbs(self) -> List[str]:
        return utility.list_collections(using=self.alias)

    def has_collection(self, name: str) -> bool:
        return utility.has_collection(name, using=self.alias)

    def drop_db(self, name: str):
        if self.has_collection(name):
            utility.drop_collection(name, using=self.alias)
        print(f"[Milvus@{self.alias}]已删除集合 {self.name}")

    def useMilvus(
        self,
        name: str,
        model_embed=None,  # 模型的文本向量化函数
        fields=None,       # 字段定义
        metric_type=None,  # 索引类型
        show: bool = True  # 是否输出日志
    ):
        from .useMilvus import useMilvus
        return useMilvus(
            client=self,
            name=name,
            model_embed=model_embed,
            fields=fields,
            metric_type=metric_type,
            show=show
        )

    def print_info(self):
        print(f"Host: {self.host}:{self.port}")
        print(f"User: {self.user or '(未设置)'}")
        print(f"Collections: {self.dbs()}")

    def drop_SQL(self):
        """清空整个数据库（删除所有集合）"""
        are_you_sure = input(
            f"[Milvus@{self.alias}] 确认清空整个数据库？(y/n) ").strip().lower()
        if are_you_sure != "y":
            print(f"[Milvus@{self.alias}] 已取消")
            return

        for coll in self.dbs():
            self.drop_db(coll)
        print("已清空数据库")

    def __repr__(self):
        ver = utility.get_server_version(using=self.alias)
        return f"<linkMilvus alias='{self.alias}' host='{self.host}:{self.port}' version='{ver}'>"
